import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim

# 定义转换操作：将图片转化为Tensor，并归一化像素值至[0, 1]
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 下载MNIST数据集并加载
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 定义一个多层感知机(MLP)用于分类任务
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        # 输入维度784（28*28），隐藏层尺寸128和64，输出类别数为10
        self.fc1 = nn.Linear(in_features=784, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=64)
        self.fc3 = nn.Linear(in_features=64, out_features=10)

    def forward(self, x):
        # 将输入x展平成二维张量 [batch_size, 784] 并通过各层计算
        x = x.view(-1, 28 * 28)  
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

model = MLP()

# 设置损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 开始训练过程...
epochs = 5
for epoch in range(epochs):
    for images, labels in train_loader:
        optimizer.zero_grad()   # 清空梯度
        
        output = model(images)     # 前向传播得到预测结果
        loss = criterion(output, labels)  # 计算损失值
    
        loss.backward()       # 反向传播求导数更新权重参数
        optimizer.step()      # 执行单步参数更新

print("完成训练!")